
from app.db.models.asset_tasks import get_monitor_urls

from . import BaseHandler
from .utils import get_warning_id
from .events import Event
from ...libs.score import Score


class VulHandler(BaseHandler):
    def handle(self):
        return self._handle_vul(task_settings=self.task_settings, response=self.response)

    def _handle_vul(self, task_settings, response):
        task_allowed_vul_types = set(task_settings.get("alert", {}).get("vulType", []))
        statistics = response.get("statistics") or {}
        statistics.update({
            "urls": len(get_monitor_urls(task_session_id=self.task.taskSessionId, is_source_data=True)),
            "scan_count": get_monitor_urls(
                task_session_id=self.task.taskSessionId, is_scan_count=True, is_source_data=True),
        })
        result = {
            "addition": {
                "statistics": statistics,
                "site_info": {"waf_list": response.get("site_info", {}).get("waf_list", [])},
                "target_ip": response.get("target_ip", "")
            },
            "warnings": []
        }
        if target_status := Event._target_status(response.get("site_info", {})):
            result["targetStatus"] = target_status
        severityMap = Score.severityMap
        for item in response.get("vuln", []):
            item_vul_type = item.get("vuln_type", [{}])[0].get("en_us")
            if item_vul_type in task_allowed_vul_types:
                affects = item.get("url") or item.get("affect")
                warning = {
                    "id": get_warning_id(
                        risk_title=item.get("source_id"),
                        task_type="vul",
                        url=affects,
                        job_id=self.task.jobId
                    ),
                    "level": severityMap.get(item["severity"], {}).get("level"),
                    "levelName": severityMap.get(item["severity"], {}).get("zh_CN"),
                    "title": item.get("name", {}).get("zh_cn"),
                    "category": item.get("vuln_type", [{}])[0].get("zh_cn"),
                    "affects": affects,
                    "detailText": item.get("extra_detail"),
                    "recommendation": item.get("recommendation", {}).get("zh_cn"),
                    "impact": item.get("impact", {}).get("zh_cn"),
                    "description": item.get("description", {}).get("zh_cn"),
                    "reference": item.get("reference"),
                    "cnvd": item.get("cnvd"),
                    "cve": item.get("cve"),
                    "cnnvd": item.get("cnnvd"),
                    "foundAt": item.get("found_at"),
                    "traffics": item.get("traffics"),
                    "payload": item.get("payload"),
                    "isResolved": False
                }
                result["warnings"].append(warning)
        return result
